package vulnerability

import (
	"github.com/gin-gonic/gin"
	"github.com/jweny/pocassist/api/msg"
	"github.com/jweny/pocassist/pkg/db"
	"github.com/jweny/pocassist/pkg/util"
	"github.com/unknwon/com"
)

type VulSerializer struct {
	// 返回给前端的字段
	WriterName		string		`json:"writer_name"`
	WebAppName		string		`json:"webapp_name"`
	Id       	int         	`json:"id"`
	NameZh 		string  		`json:"name_zh"`
	Cve 		string     		`json:"cve"`
	Cnnvd 		string   		`json:"cnnvd"`
	Severity	string    		`json:"severity"`
	Category 	string    		`json:"category"`
	Description string    		`json:"description"`
	Suggestion	string  		`json:"suggestion"`
	Language 	string    		`json:"language"`
	Webapp		int      		`json:"webapp"`
}

// @Summary vul detail
// @Tags Vul
// @Description 详情
// @Produce  json
// @Security token
// @Param id path int true "ID"
// @Success 200 {object} msg.Response
// @Failure 200 {object} msg.Response
// @Router /api/v1/vul/{id}/ [get]
func Detail(c *gin.Context) {
	id := com.StrTo(c.Param("id")).MustInt()
	var data interface {}
	if db.ExistVulnerabilityByID(id) {
		data = db.GetVulnerability(id)
		c.JSON(msg.SuccessResp(data))
		return
	} else {
		c.JSON(msg.ErrResp("record not found"))
		return
	}
}

// @Summary vul list
// @Tags Vul
// @Description 列表
// @Produce  json
// @Security token
// @Param page query int true "Page"
// @Param pagesize query int true "Pagesize"
// @Param field query db.VulnerabilitySearchField false "field"
// @Success 200 {object} msg.Response
// @Failure 200 {object} msg.Response
// @Router /api/v1/vul/ [get]
func Get(c *gin.Context) {
	data := make(map[string]interface{})
	field := db.VulnerabilitySearchField{
		Search:"",
		WebappField:-1,
		CategoryField:"",}

	// 分页
	page, _ := com.StrTo(c.Query("page")).Int()
	pageSize, _ := com.StrTo(c.Query("pagesize")).Int()

	// 查询条件
	if arg := c.Query("search"); arg != "" {
		field.Search = arg
	}
	if arg := c.Query("productField"); arg != "" {
		appId := com.StrTo(arg).MustInt()
		field.WebappField = appId
	}
	if arg := c.Query("typeField"); arg != "" {
		field.CategoryField = arg
	}

	vuls := db.GetVulnerabilities(page, pageSize, &field)
	var vulRespData []VulSerializer
	// 获取上传者
	token := c.Request.Header.Get("Authorization")
	claims, _ := util.ParseToken(token)

	for _, vul := range vuls {
		var appName string
		if vul.ForeignWebapp != nil {
			appName = vul.ForeignWebapp.Name
		} else {
			appName = ""
		}
		vulRespData = append(vulRespData, VulSerializer{
			WriterName:    claims.Username,
			WebAppName:    appName,
			Id:vul.Id,
			NameZh:vul.NameZh,
			Cve:vul.Cve,
			Cnnvd:vul.Cnnvd,
			Severity:vul.Severity,
			Category:vul.Category,
			Description:vul.Description,
			Suggestion:vul.Suggestion,
			Language:vul.Language,
			Webapp:vul.Webapp,
		})
	}
	data["data"] = vulRespData
	total := db.GetVulnerabilitiesTotal(&field)
	data["total"] = total
	c.JSON(msg.SuccessResp(data))
	return
}

// 给swagger用的结构体
type SwagVulnerability struct {
	Id       	int         `gorm:"primary_key" json:"id"`
	NameZh 		string  `gorm:"column:name_zh" json:"name_zh" binding:"required"`
	Cve 		string     `gorm:"column:cve" json:"cve"`
	Cnnvd 		string   `gorm:"column:cnnvd" json:"cnnvd"`
	Severity	string    `gorm:"column:severity" json:"severity"`
	Category 	string    `gorm:"column:category" json:"category"`
	Description string    `gorm:"type:longtext" json:"description"`
	Suggestion	string  `gorm:"type:longtext" json:"suggestion"`
	Language 	string    `gorm:"column:language" json:"language"`
	Webapp		int      `gorm:"column:webapp" json:"webapp"`
	ForeignWebapp *db.Webapp `gorm:"foreignKey:Webapp"`
}

// @Summary vul add
// @Tags Vul
// @Description 新增
// @Produce  json
// @Security token
// @Param vul body SwagVulnerability true "vul"
// @Success 200 {object} msg.Response
// @Failure 200 {object} msg.Response
// @Router /api/v1/vul/ [post]
func Add(c *gin.Context) {
	vul := db.Vulnerability{}
	err := c.ShouldBindJSON(&vul)
	if err != nil {
		c.JSON(msg.ErrResp("漏洞名称不可为空"))
		return
	}
	if db.ExistVulnerabilityByNameZh(vul.NameZh){
		c.JSON(msg.ErrResp("漏洞名称已存在"))
		return
	} else {
		db.AddVulnerability(vul)
		c.JSON(msg.SuccessResp(vul))
		return
	}
}

// @Summary vul update
// @Tags Vul
// @Description 更新
// @Produce  json
// @Security token
// @Param id path int true "ID"
// @Param vul body SwagVulnerability true "vul"
// @Success 200 {object} msg.Response
// @Failure 200 {object} msg.Response
// @Router /api/v1/vul/{id}/ [put]
func Update(c *gin.Context) {
	id := com.StrTo(c.Param("id")).MustInt()
	vul := db.Vulnerability{}
	err := c.ShouldBindJSON(&vul)
	if err != nil {
		c.JSON(msg.ErrResp("漏洞名称不可为空"))
		return
	}

	if db.ExistVulnerabilityByID(id){
		db.EditVulnerability(id, vul)
		c.JSON(msg.SuccessResp(vul))
	} else {
		c.JSON(msg.ErrResp("record not found"))
		return
	}
}

// @Summary vul delete
// @Tags Vul
// @Description 删除
// @Produce  json
// @Security token
// @Param id path int true "ID"
// @Success 200 {object} msg.Response
// @Failure 200 {object} msg.Response
// @Router /api/v1/vul/{id}/ [delete]
func Delete(c *gin.Context) {
	id := com.StrTo(c.Param("id")).MustInt()

	if db.ExistVulnerabilityByID(id) {
		db.DeleteVulnerability(id)
		c.JSON(msg.SuccessResp("删除成功"))
		return
	} else {
		c.JSON(msg.ErrResp("record not found"))
		return
	}

}

type BasicObj struct {
	Name 	string 	`json:"name"`
	Label 	string	`json:"label"`
}

// @Summary vul basic
// @Tags Vul
// @Description 前端需要的基础信息
// @Produce  json
// @Security token
// @Success 200 {object} msg.Response
// @Failure 200 {object} msg.Response
// @Router /api/v1/vul/basic/ [get]
func Basic(c *gin.Context) {
	var LanguageChoice []BasicObj
	for _, v := range []string{"Any","ASP","JAVA","Python","GO","NodeJS","PHP","Ruby","ASPX",} {
		LanguageChoice = append(LanguageChoice, BasicObj{Name: v, Label:v})
	}
	var AffectChoice []BasicObj
	for _, v := range []string{"server","text","directory","url","appendparam","replaceparam","script"} {
		AffectChoice = append(AffectChoice, BasicObj{Name: v, Label:v})
	}
	var LevelChoice []BasicObj
	for _, v := range []string{"high","middle","low","info",} {
		LevelChoice = append(LevelChoice, BasicObj{Name: v, Label:v})
	}
	var TypeChoice []BasicObj
	for _, v := range []string{"SQL 注入","命令执行","信息泄漏","其他类型","目录穿越","未授权","任意文件下载","任意文件读取","反序列化","任意文件写入","弱口令","权限提升","文件上传","远程文件包含","跨站请求伪造","跨站脚本XSS","XPath注入","XML注入","服务器端请求伪造","解析错误","本地文件包含","配置错误"} {
		TypeChoice = append(TypeChoice, BasicObj{Name: v, Label:v})
	}

	data := make(map[string]interface{})
	data["VulLanguage"] = LanguageChoice
	data["VulLevel"] = LevelChoice
	data["ModuleAffects"] = AffectChoice
	data["VulType"] = TypeChoice

	c.JSON(msg.SuccessResp(data))
	return
}